#!/usr/bin/env python
# -*- coding: utf-8 -*-


"""
Created on 2018年8月4日
@author: Irony
@site: https://pyqt.site , https://github.com/PyQt5
@email: 892768447@qq.com
@file: QListView.显示自定义Widget
@description:
"""

try:
    from PyQt5.QtCore import QSize, pyqtSignal
    from PyQt5.QtGui import QStandardItemModel, QStandardItem, QImage, QPixmap
    from PyQt5.QtWidgets import QMessageBox, QFileDialog, QGridLayout, QDialog, QListView, QWidget, QHBoxLayout, QLineEdit, \
        QPushButton, QApplication, QLabel, QSpinBox
except ImportError:
    from PySide2.QtCore import QSize
    from PySide2.QtGui import QStandardItemModel, QStandardItem
    from PySide2.QtWidgets import QListView, QDialog, QWidget, QHBoxLayout, QLineEdit, \
        QPushButton, QApplication

from ui.CameraCollect import Ui_Dialog as CameraCollect

import cv2
import numpy as np
import threading
import time
import os


class CameraCollectDialog(QDialog, CameraCollect):

    cameras = []
    signal = pyqtSignal(QLabel, np.ndarray)
    signal_save = pyqtSignal(str, np.ndarray)
    signal_text = pyqtSignal(str)

    def __init__(self, *args, **kwargs):
        super(CameraCollectDialog, self).__init__(*args, **kwargs)
        self.setupUi(self)
        self.setWindowTitle("红外图像采集工具")

        self.pushButton_scan.clicked.connect(self.scanCamera)
        self.pushButton_start.clicked.connect(self.startCamera)
        self.pushButton_stop.clicked.connect(self.stopCamera)
        self.pushButton_select_dir.clicked.connect(self.selectDir)
        self.lineEdit_dir.setEnabled(False)

        self.signal.connect(self.cameraShow)
        self.signal_save.connect(self.saveImage)
        self.signal_text.connect(self.textAppend)

    def selectDir(self):
        res = QFileDialog.getExistingDirectory(self, "选择保存路径", "./")
        if res != '':
            self.lineEdit_dir.setText(res)

    def createDir(self, path):
        isExists = os.path.exists(path)
        if not isExists:
            os.makedirs(path)
            return True
        else:
            # 如果目录存在则不创建，并提示目录已存在
            return False

    def scanCamera(self):
        if self.cameras:
            print(self.cameras)
            for _ in range(len(self.cameras)):
                c = self.cameras[0]
                self.gridLayout_camera.removeWidget(c['label'])
                self.cameras.pop(0)
        print(self.cameras)
        cnt = 0
        for device in range(0, 5):
            stream = cv2.VideoCapture(device)
            grabbed = stream.grab()
            stream.release()
            if not grabbed:
                break

            label = QLabel("%s" % cnt)
            self.textBrowser.append('搜索到摄像头:' + str(device))
            self.cameras.append({'name': cnt, 'label': label})
            self.gridLayout_camera.addWidget(label, int(cnt // 2), int(cnt // 2) + (cnt % 2), 1, 1)

            cnt = cnt + 1
        self.textBrowser.append('搜索完成')
        return cnt

    def startCamera(self):
        inv = self.spinBox_time.value()
        save_dir = self.lineEdit_dir.text()
        if save_dir == '':
            QMessageBox.warning(self, "提示", "请选择保存路径！", QMessageBox.Yes)
            return
        self.threads = []
        self.threads_run = True
        for c in self.cameras:
            path = '%s/%s' % (save_dir, c['name'])
            self.createDir(path)
            camera_thread = threading.Thread(target=self.cameraThread, args=[c['name'], c['label'], path, inv])
            camera_thread.setDaemon(True)
            camera_thread.start()
            self.threads.append(camera_thread)

        self.pushButton_start.setEnabled(False)
        self.pushButton_select_dir.setEnabled(False)
        self.spinBox_time.setEnabled(False)
        self.pushButton_scan.setEnabled(False)
        self.pushButton_stop.setEnabled(True)

    def stopCamera(self):
        self.threads_run = False
        self.threads = []
        self.pushButton_start.setEnabled(True)
        self.spinBox_time.setEnabled(True)
        self.pushButton_select_dir.setEnabled(True)
        self.pushButton_scan.setEnabled(True)
        self.pushButton_stop.setEnabled(False)

    def cameraThread(self, name, label, path, inv):
        self.signal_text.emit("正在开启摄像头:" + str(name))
        cnt = 0
        t0 = time.time()
        capture = cv2.VideoCapture(name)
        while self.threads_run:
            try:
                if capture:
                    _, img = capture.read()
                    self.signal.emit(label, img)
                    now = time.time()
                    if (now - t0) >= inv:
                        cnt += 1
                        t0 = now
                        file = '%s/%05d.bmp' % (path, cnt)
                        cv2.imwrite(file, img)
                        self.signal_text.emit('save:' + file)
                        if cnt > 10000:
                            self.stopCamera()
            except Exception as e:
                print(e)

    def cameraShow(self, label, img):
        # 提取图像的通道和尺寸，用于将OpenCV下的image转换成Qimage
        if len(img.shape) == 3:
            height, width, channel = img.shape
            bytesPerline = 3 * width
            QImg = QImage(img.data, width, height, bytesPerline,
                          QImage.Format_RGB888).rgbSwapped()
        elif len(img.shape) == 2:  # 灰度图
            height, width = img.shape
            bytesPerline = width
            QImg = QImage(img.data, width, height, bytesPerline,
                          QImage.Format_Indexed8)
        else:
            return
        # 将QImage显示出来
        label.setPixmap(QPixmap.fromImage(QImg))

    def saveImage(self, path, img):
        cv2.imwrite(path, img)
        self.textBrowser.append('save: %s' % path)

    def textAppend(self, info):
        self.textBrowser.append(info)


if __name__ == '__main__':
    import sys

    app = QApplication(sys.argv)
    w = CameraCollectDialog()
    w.show()
    sys.exit(app.exec_())
